{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于with_cache_convert函数：\n",
    "1. 首先，我们定义了一个函数`with_cache_convert`，它接收一个参数`func`，即需要被装饰的函数。\n",
    "2. 在函数中，我们首先检查缓存中是否已经存在了该函数的输出结果。如果已经存在，则直接返回缓存中的结果。\n",
    "3. 如果缓存中不存在该函数的输出结果，则调用`func`函数，并将结果存储在缓存中。\n",
    "4. 最后，返回函数的输出结果。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "dynamic_axes 参数\n",
    "作用：\n",
    "dynamic_axes 用于指定模型的输入和输出张量中哪些维度是动态的（即可以在运行时改变大小）。这对于支持可变批次大小或可变序列长度的模型非常重要。\n",
    "\n",
    "参数结构：\n",
    "dynamic_axes 是一个字典，键是输入或输出的名称，值是一个字典，指定该张量的哪些维度是动态的。例如：\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于onnx_test_with_cachehanshu：\n",
    "这个函数的功能是使用缓存机制将onnx模型转换为onnxruntime的推理会话。这个函数使用了装饰器@functools.wraps(func)，这个装饰器会将被装饰的函数的元数据（如函数名、文档字符串等）传递给装饰器函数。\n",
    "\n",
    "在函数内部，首先检查缓存中是否已经存在与onnx文件路径相同的onnxruntime会话。如果存在，则直接返回该会话。如果不存在，则使用onnxruntime的InferenceSession类创建一个新的onnxruntime会话，"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional, Tuple\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "import onnxruntime as ort\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "class Config:\n",
    "    def __init__(self, hidden_size: int = 512, num_attention_heads: int = 8) -> None:\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_attention_heads = num_attention_heads\n",
    "\n",
    "\n",
    "class AttentionWithCache(nn.Module):\n",
    "    def __init__(self, config: Config) -> None:\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.hidden_size = config.hidden_size\n",
    "        self.num_heads = config.num_attention_heads\n",
    "        self.head_dim = self.hidden_size // self.num_heads\n",
    "        assert self.head_dim * self.num_heads == self.hidden_size, \"hidden_size must be divisible by num_heads.\"\n",
    "        \n",
    "        self.query_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)\n",
    "        self.key_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)\n",
    "        self.value_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)\n",
    "        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)\n",
    "\n",
    "    def forward(self, hidden_states: torch.Tensor, key_cache: Optional[torch.Tensor] = None, value_cache: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        bsz, q_len, _ = hidden_states.size()\n",
    "\n",
    "        # 计算 query\n",
    "        query_states = self.query_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 如果存在缓存，使用缓存并计算新的 key 和 value，然后拼接它们\n",
    "        if key_cache is not None:\n",
    "            key_states = torch.cat([key_cache, self.key_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)], dim=2)\n",
    "        else:\n",
    "            key_states = self.key_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        if value_cache is not None:\n",
    "            value_states = torch.cat([value_cache, self.value_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)], dim=2)\n",
    "        else:\n",
    "            value_states = self.value_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) / math.sqrt(self.head_dim)\n",
    "        attn_weights = torch.softmax(attn_weights, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        attn_output = torch.matmul(attn_weights, value_states).transpose(1, 2).reshape(bsz, q_len, self.hidden_size)\n",
    "        attn_output = self.o_proj(attn_output)\n",
    "\n",
    "        # 返回注意力输出以及更新后的 key 和 value 缓存\n",
    "        return attn_output, key_states, value_states\n",
    "\n",
    "\n",
    "def with_cache_convert(model: AttentionWithCache, save_path='model.onnx') -> None:\n",
    "    # 将模型设置为评估模式，关闭 dropout 和 batch normalization 的训练行为\n",
    "    model.eval()\n",
    "\n",
    "    # 创建随机输入张量，形状为 [batch_size, sequence_length, hidden_size]\n",
    "    # batch_size=1, sequence_length=10, hidden_size 从模型中获取\n",
    "    dummy_input = torch.randn(1, 10, model.hidden_size).to(device)\n",
    "\n",
    "    # 创建随机的 key 缓存张量，形状为 [batch_size, num_heads, sequence_length, head_dim]\n",
    "    # batch_size=1, num_heads 从模型中获取, sequence_length=10, head_dim 从模型中获取\n",
    "    key_cache = torch.randn(1, model.num_heads, 10, model.head_dim).to(device)\n",
    "\n",
    "    # 创建随机的 value 缓存张量，形状为 [batch_size, num_heads, sequence_length, head_dim]\n",
    "    # batch_size=1, num_heads 从模型中获取, sequence_length=10, head_dim 从模型中获取\n",
    "    value_cache = torch.randn(1, model.num_heads, 10, model.head_dim).to(device)\n",
    "\n",
    "    # 导出模型为 ONNX 格式\n",
    "    torch.onnx.export(\n",
    "        model,  # 要导出的模型\n",
    "        (dummy_input, key_cache, value_cache),  # 模型的输入，包括隐藏状态和缓存\n",
    "        save_path,  # 导出的 ONNX 文件保存路径\n",
    "        input_names=['hidden_states', 'key_cache', 'value_cache'],  # 输入的名称\n",
    "        output_names=['output', 'updated_key_cache', 'updated_value_cache'],  # 输出的名称\n",
    "        dynamic_axes={\n",
    "            # 指定动态维度：\n",
    "            # - hidden_states: 第 0 维（batch_size）和第 1 维（sequence_length）是动态的\n",
    "            'hidden_states': {0: 'batch_size', 1: 'sequence_length'},\n",
    "            # - key_cache: 第 0 维（batch_size）和第 2 维（sequence_length）是动态的\n",
    "            'key_cache': {0: 'batch_size', 2: 'sequence_length'},\n",
    "            # - value_cache: 第 0 维（batch_size）和第 2 维（sequence_length）是动态的\n",
    "            'value_cache': {0: 'batch_size', 2: 'sequence_length'},\n",
    "            # - output: 第 0 维（batch_size）和第 1 维（sequence_length）是动态的\n",
    "            'output': {0: 'batch_size', 1: 'sequence_length'},\n",
    "            # - updated_key_cache: 第 0 维（batch_size）和第 2 维（sequence_length）是动态的\n",
    "            'updated_key_cache': {0: 'batch_size', 2: 'sequence_length'},\n",
    "            # - updated_value_cache: 第 0 维（batch_size）和第 2 维（sequence_length）是动态的\n",
    "            'updated_value_cache': {0: 'batch_size', 2: 'sequence_length'}\n",
    "        },\n",
    "        opset_version=11  # 使用 ONNX 算子集的第 11 版\n",
    "    )\n",
    "\n",
    "def onnx_test_with_cache(hidden_states: np.ndarray, key_cache: Optional[np.ndarray] = None, value_cache: Optional[np.ndarray] = None, model_path: str='model.onnx') -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "    # 初始化 ONNX 运行时会话，加载指定路径的 ONNX 模型\n",
    "    ort_session = ort.InferenceSession(model_path)\n",
    "\n",
    "    # 获取模型的配置信息：\n",
    "    # - hidden_size 从输入 hidden_states 的最后一维推断\n",
    "    # - num_attention_heads 默认为 8\n",
    "    config = Config(hidden_size=hidden_states.shape[-1], num_attention_heads=8)\n",
    "\n",
    "    # 计算每个注意力头的维度\n",
    "    num_heads = config.num_attention_heads\n",
    "    head_dim = config.hidden_size // num_heads   \n",
    "\n",
    "    # 如果 key_cache 未提供，则初始化为零矩阵，形状为 [batch_size, num_heads, 0, head_dim]\n",
    "    if key_cache is None:\n",
    "        key_cache = np.zeros((hidden_states.shape[0], num_heads, 0, head_dim), dtype=np.float32)\n",
    "\n",
    "    # 如果 value_cache 未提供，则初始化为零矩阵，形状为 [batch_size, num_heads, 0, head_dim]\n",
    "    if value_cache is None:\n",
    "        value_cache = np.zeros((hidden_states.shape[0], num_heads, 0, head_dim), dtype=np.float32)\n",
    "\n",
    "    # 准备 ONNX 运行时的输入数据\n",
    "    ort_inputs = {\n",
    "        'hidden_states': hidden_states,  # 输入的隐藏状态\n",
    "        'key_cache': key_cache,  # 输入的 key 缓存\n",
    "        'value_cache': value_cache  # 输入的 value 缓存\n",
    "    }\n",
    "\n",
    "    # 运行 ONNX 模型，获取输出\n",
    "    ort_outs = ort_session.run(None, ort_inputs)\n",
    "\n",
    "    # 返回输出和更新后的缓存：\n",
    "    # - ort_outs[0]: 模型输出\n",
    "    # - ort_outs[1]: 更新后的 key 缓存\n",
    "    # - ort_outs[2]: 更新后的 value 缓存\n",
    "    return ort_outs[0], ort_outs[1], ort_outs[2]\n",
    "\n",
    "\n",
    "def main(torch_model):\n",
    "    config = Config(hidden_size=512, num_attention_heads=8)\n",
    "    hidden_size = config.hidden_size\n",
    "    attention_layer_with_cache = AttentionWithCache(config).to(device)\n",
    "    attention_layer_with_cache.load_state_dict(torch.load(torch_model))\n",
    "    with_cache_convert(attention_layer_with_cache)\n",
    "    hidden_states=np.random.randn(1, 10, hidden_size).astype(np.float32)\n",
    "    output, key_cache, value_cache = onnx_test_with_cache(hidden_states, model_path='model.onnx')\n",
    "    print(output.shape, key_cache.shape, value_cache.shape)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main('model.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorchgpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
